#include "YOLODetector.h"
#include <iostream>
#include <chrono>
#include <string>

/**
 * @brief 显示程序使用帮助信息
 * @param programName 程序名称
 */
void printUsage(const char* programName) {
    std::cout << "用法: " << programName << " [选项]" << std::endl;
    std::cout << "选项:" << std::endl;
    std::cout << "  -m <path>   指定模型文件路径 (默认: best.onnx)" << std::endl;
    std::cout << "  -i <path>   指定输入图像路径 (默认: 1.jpg)" << std::endl;
    std::cout << "  -c <float>  指定置信度阈值 (默认: 0.25)" << std::endl;
    std::cout << "  -n <float>  指定NMS阈值 (默认: 0.45)" << std::endl;
    std::cout << "  -h          显示帮助信息" << std::endl;
}

/**
 * @brief 主函数
 * @param argc 参数数量
 * @param argv 参数数组
 * @return 程序退出状态码
 */
int main(int argc, char** argv) {
    // 默认参数
    std::string modelPath = "/home/fourzkw/Project/yolo_test/yolo_test_cpp/best.onnx";
    std::string imagePath = "/home/fourzkw/Project/yolo_test/yolo_test_cpp/1.jpg";
    float confThreshold = 0.25f;
    float nmsThreshold = 0.45f;

    // 解析命令行参数
    for (int i = 1; i < argc; i++) {
        std::string arg = argv[i];
        
        if (arg == "-h") {
            printUsage(argv[0]);
            return 0;
        } else if (arg == "-m" && i + 1 < argc) {
            modelPath = argv[++i];
        } else if (arg == "-i" && i + 1 < argc) {
            imagePath = argv[++i];
        } else if (arg == "-c" && i + 1 < argc) {
            confThreshold = std::stof(argv[++i]);
        } else if (arg == "-n" && i + 1 < argc) {
            nmsThreshold = std::stof(argv[++i]);
        }
    }

    // 显示程序参数
    std::cout << "使用参数:" << std::endl;
    std::cout << "- 模型: " << modelPath << std::endl;
    std::cout << "- 图像: " << imagePath << std::endl;
    std::cout << "- 置信度阈值: " << confThreshold << std::endl;
    std::cout << "- NMS阈值: " << nmsThreshold << std::endl;

    // 创建YOLO检测器实例
    YOLODetector detector;

    // 初始化模型
    if (!detector.initialize(modelPath)) {
        std::cerr << "模型初始化失败!" << std::endl;
        return -1;
    }

    // 读取输入图像
    cv::Mat image = cv::imread(imagePath);
    if (image.empty()) {
        std::cerr << "无法读取图像: " << imagePath << std::endl;
        return -1;
    }

    // 检测开始计时
    auto start = std::chrono::high_resolution_clock::now();

    // 执行目标检测
    std::vector<Detection> detections = detector.detect(image, confThreshold, nmsThreshold);

    // 检测结束计时
    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> diff = end - start;

    // 显示检测结果统计
    std::cout << "检测到 " << detections.size() << " 个目标，用时 " 
              << diff.count() * 1000 << " 毫秒" << std::endl;

    // 在原图上绘制检测结果
    detector.drawDetections(image, detections);

    // 创建窗口并显示结果
    cv::namedWindow("检测结果", cv::WINDOW_NORMAL);
    cv::imshow("检测结果", image);
    
    // 保存结果图像
    std::string outputPath = "result.jpg";
    cv::imwrite(outputPath, image);
    std::cout << "结果已保存至: " << outputPath << std::endl;

    // 等待用户按键
    cv::waitKey(0);
    
    return 0;
} 